// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#include <benchmark/benchmark.h>

#include <nebula/util/key_value_metadata.h>
#include <nebula/util/string.h>

#include <nebula/compute/row/grouper.h>
#include <nebula/testing/ktest_util.h>
#include <nebula/testing/random.h>
#include <nebula/util/benchmark_util.h>

namespace nebula {
namespace compute {

constexpr auto kSeed = 0x0ff1ce;
constexpr int64_t kRound = 16;
constexpr double true_and_unique_probability = 0.2;

static ExecBatch MakeRandomExecBatch(const DataTypeVector& types, int64_t num_rows,
                                     double null_probability,
                                     int64_t alignment = kDefaultBufferAlignment,
                                     MemoryPool* memory_pool = nullptr) {
  random::RandomArrayGenerator rng(kSeed);
  auto num_types = static_cast<int>(types.size());

  // clang-format off
  // For unique probability:
  // The proportion of Unique determines the number of groups.
  // 1. In most scenarios, unique has a small proportion and exists
  // 2. In GroupBy/HashJoin are sometimes used for deduplication and
  // in that use case the key is mostly unique
  auto metadata = key_value_metadata(
      {
        "null_probability",
        "true_probability",  // for boolean type
        "unique"             // for string type
      },
      {
          internal::ToChars(null_probability),
          internal::ToChars(true_and_unique_probability),
          internal::ToChars(static_cast<int32_t>(num_rows *
                            true_and_unique_probability))
      });
  // clang-format on

  std::vector<Datum> values;
  values.resize(num_types);
  for (int i = 0; i < num_types; ++i) {
    auto field = ::nebula::field("", types[i], metadata);
    values[i] = rng.ArrayOf(*field, num_rows, alignment, memory_pool);
  }

  return ExecBatch(std::move(values), num_rows);
}

static void GrouperBenchmark(benchmark::State& state, const ExecSpan& span,
                             ExecContext* ctx = nullptr) {
  uint32_t num_groups = 0;
  for (auto _ : state) {
    ASSIGN_OR_ABORT(auto grouper, Grouper::create(span.get_types(), ctx));
    for (int i = 0; i < kRound; ++i) {
      ASSIGN_OR_ABORT(auto group_ids, grouper->consume(span));
    }
    num_groups = grouper->num_groups();
  }

  state.SetItemsProcessed(state.iterations() * kRound * span.length);
  state.counters["num_groups"] = num_groups;
  state.counters["uniqueness"] = static_cast<double>(num_groups) / (kRound * span.length);
}

static void GrouperWithMultiTypes(benchmark::State& state, const DataTypeVector& types) {
  auto ctx = default_exec_context();

  RegressionArgs args(state, false);
  const int64_t num_rows = args.size;
  const double null_proportion = args.null_proportion;

  auto exec_batch = MakeRandomExecBatch(types, num_rows, null_proportion,
                                        kDefaultBufferAlignment, ctx->memory_pool());
  ExecSpan exec_span(exec_batch);
  ASSIGN_OR_ABORT(auto grouper, Grouper::create(exec_span.get_types(), ctx));
  GrouperBenchmark(state, exec_span, ctx);
}

void SetArgs(benchmark::internal::Benchmark* bench) {
  BenchmarkSetArgsWithSizes(bench, {1 << 10, 1 << 12});
}

// basic types
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{boolean}", {boolean()})->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32}", {int32()})->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int64}", {int64()})->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{utf8}", {utf8()})->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{fixed_size_binary(32)}",
                  {fixed_size_binary(32)})
    ->Apply(SetArgs);

// combination types
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{boolean, utf8}", {boolean(), utf8()})
    ->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, int32}", {int32(), int32()})
    ->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int64, int32}", {int64(), int32()})
    ->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{boolean, int64, utf8}",
                  {boolean(), int64(), utf8()})
    ->Apply(SetArgs);

// combination types requiring column resorted
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, boolean, utf8}",
                  {int32(), boolean(), utf8()})
    ->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes, "{int32, int64, boolean, utf8}",
                  {int32(), int64(), boolean(), utf8()})
    ->Apply(SetArgs);
BENCHMARK_CAPTURE(GrouperWithMultiTypes,
                  "{utf8, int32, int64, fixed_size_binary(32), boolean}",
                  {utf8(), int32(), int64(), fixed_size_binary(32), boolean()})
    ->Apply(SetArgs);

}  // namespace compute
}  // namespace nebula
